"""
基于量子计算的资产组合优化策略(Zipline实现)
策略特点：
1. 使用经典优化算法模拟量子组合优化
2. 支持大规模资产组合的快速优化
3. 实现定期再平衡
"""

from zipline.api import order_target_percent, record, symbol, get_datetime
from zipline.finance import commission, slippage
import numpy as np
from scipy.optimize import minimize

def initialize(context):
    # 策略参数
    context.assets = [symbol('AAPL'), symbol('MSFT'), symbol('GOOG'), symbol('AMZN')]
    context.risk_aversion = 0.5
    context.rebalance_frequency = 21  # 每月再平衡 (21个交易日)
    context.last_rebalance = None
    
    # 设置交易成本
    context.set_commission(commission.PerShare(cost=0.001, min_trade_cost=1))
    context.set_slippage(slippage.FixedSlippage(spread=0.01))
    
    # 初始化状态
    context.weights = np.ones(len(context.assets)) / len(context.assets)  # 初始等权重

def portfolio_objective(weights, returns, cov_matrix, risk_aversion):
    """组合优化目标函数"""
    portfolio_return = np.sum(returns * weights)
    portfolio_risk = np.sqrt(weights.T @ cov_matrix @ weights)
    return - (portfolio_return - risk_aversion * portfolio_risk)

def optimize_portfolio(returns, cov_matrix, risk_aversion):
    """经典优化模拟量子组合优化"""
    n_assets = len(returns)
    initial_weights = np.ones(n_assets) / n_assets
    
    # 约束条件
    constraints = ({'type': 'eq', 'fun': lambda w: np.sum(w) - 1})  # 权重和为1
    bounds = [(0, 1) for _ in range(n_assets)]  # 不允许做空
    
    # 优化
    result = minimize(portfolio_objective, initial_weights, 
                     args=(returns, cov_matrix, risk_aversion),
                     method='SLSQP', bounds=bounds, constraints=constraints)
    
    return result.x

def handle_data(context, data):
    # 检查是否需要再平衡
    current_date = get_datetime().date()
    if context.last_rebalance is None or (current_date - context.last_rebalance).days >= context.rebalance_frequency:
        # 计算历史收益和协方差
        prices = data.history(context.assets, 'price', 252, '1d')  # 1年历史数据
        returns = prices.pct_change().dropna()
        
        # 计算预期收益和协方差矩阵
        expected_returns = returns.mean().values
        cov_matrix = returns.cov().values
        
        # 优化组合
        context.weights = optimize_portfolio(expected_returns, cov_matrix, context.risk_aversion)
        context.last_rebalance = current_date
        
        # 执行再平衡
        for i, asset in enumerate(context.assets):
            order_target_percent(asset, context.weights[i])
    
    # 记录状态
    portfolio_value = context.portfolio.portfolio_value
    positions = [context.portfolio.positions[asset].amount * data.current(asset, 'price') / portfolio_value
                 for asset in context.assets]
    
    record(positions=positions, 
          weights=context.weights, 
          portfolio_value=portfolio_value)